{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tensor and Add Operation\n",
    "\n",
    "ttnn.Tensor is the central type of ttnn.\n",
    "\n",
    "It is similar to torch.Tensor in the sense that it represents multi-dimensional matrix containing elements of a single data type.\n",
    "\n",
    "The are a few key differences:\n",
    "\n",
    "- ttnn.Tensor can be stored in the SRAM or DRAM of Tenstorrent devices\n",
    "- ttnn.Tensor doesn't have a concept of the strides, however it has a concept of row-major and tile layout\n",
    "- ttnn.Tensor has support for data types not supported by torch such as `bfp8` for example\n",
    "- ttnn.Tensor's shape stores the padding added to the tensor due to TILE_LAYOUT"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import the library"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-07-02 22:32:56.948 | DEBUG    | ttnn:<module>:83 - Initial ttnn.CONFIG:\n",
      "Config{cache_path=/home/ubuntu/.cache/ttnn,model_cache_path=/home/ubuntu/.cache/ttnn/models,tmp_dir=/tmp/ttnn,enable_model_cache=false,enable_fast_runtime_mode=true,throw_exception_on_fallback=false,enable_logging=false,enable_graph_report=false,enable_detailed_buffer_report=false,enable_detailed_tensor_report=false,enable_comparison_mode=false,comparison_mode_should_raise_exception=false,comparison_mode_pcc=0.9999,root_report_path=generated/ttnn/reports,report_name=std::nullopt,std::nullopt}\n"
     ]
    }
   ],
   "source": [
    "import ttnn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Open the device\n",
    "\n",
    "Use `ttnn.open_device` to get a handle to the device and create tensors on it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2025-07-02 22:32:57.804 | info     |   SiliconDriver | Opened PCI device 0; KMD version: 2.0.0; API: 2; IOMMU: disabled (pci_device.cpp:198)\n",
      "2025-07-02 22:32:57.806 | info     |   SiliconDriver | Opened PCI device 0; KMD version: 2.0.0; API: 2; IOMMU: disabled (pci_device.cpp:198)\n",
      "2025-07-02 22:32:57.818 | info     |          Device | Opening user mode device driver (tt_cluster.cpp:174)\n",
      "2025-07-02 22:32:57.818 | info     |   SiliconDriver | Opened PCI device 0; KMD version: 2.0.0; API: 2; IOMMU: disabled (pci_device.cpp:198)\n",
      "2025-07-02 22:32:57.819 | info     |   SiliconDriver | Opened PCI device 0; KMD version: 2.0.0; API: 2; IOMMU: disabled (pci_device.cpp:198)\n",
      "2025-07-02 22:32:57.823 | info     |   SiliconDriver | Opened PCI device 0; KMD version: 2.0.0; API: 2; IOMMU: disabled (pci_device.cpp:198)\n",
      "2025-07-02 22:32:57.823 | info     |   SiliconDriver | Opened PCI device 0; KMD version: 2.0.0; API: 2; IOMMU: disabled (pci_device.cpp:198)\n",
      "2025-07-02 22:32:57.827 | info     |   SiliconDriver | Harvesting mask for chip 0 is 0x40 (NOC0: 0x40, simulated harvesting mask: 0x0). (cluster.cpp:282)\n",
      "2025-07-02 22:32:57.860 | info     |   SiliconDriver | Opened PCI device 0; KMD version: 2.0.0; API: 2; IOMMU: disabled (pci_device.cpp:198)\n",
      "2025-07-02 22:32:57.926 | info     |   SiliconDriver | Opening local chip ids/pci ids: {0}/[0] and remote chip ids {} (cluster.cpp:147)\n",
      "2025-07-02 22:32:57.930 | info     |   SiliconDriver | Software version 6.0.0, Ethernet FW version 6.15.0 (Device 0) (cluster.cpp:1039)\n",
      "2025-07-02 22:32:57.970 | info     |           Metal | AI CLK for device 0 is:   1000 MHz (metal_context.cpp:122)\n",
      "2025-07-02 22:32:58.500 | info     |           Metal | Initializing device 0. Program cache is enabled (device.cpp:429)\n"
     ]
    }
   ],
   "source": [
    "device_id = 0\n",
    "device = ttnn.open_device(device_id=device_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And now let's create a ttnn Tensor directly on the device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "shape: Shape([3, 4])\n",
      "layout: Layout.TILE\n",
      "dtype: DataType.BFLOAT16\n"
     ]
    }
   ],
   "source": [
    "ttnn_tensor = ttnn.rand((3, 4), device=device)\n",
    "\n",
    "print(f\"shape: {ttnn_tensor.shape}\")\n",
    "print(f\"layout: {ttnn_tensor.layout}\")\n",
    "print(f\"dtype: {ttnn_tensor.dtype}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As expected we get a tensor of shape [3, 4] in tile layout with a data type of bfloat16."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Type\n",
    "\n",
    "The data type of the ttnn tensor can be explicitly controlled by passing the dtype parameter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ttnn_tensor.dtype: DataType.FLOAT32\n"
     ]
    }
   ],
   "source": [
    "ttnn_tensor = ttnn.rand((3, 4), dtype=ttnn.float32, device=device)\n",
    "print(f\"ttnn_tensor.dtype: {ttnn_tensor.dtype}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Layout\n",
    "\n",
    "Tenstorrent hardware is most efficiently utilized when running tensors using [tile layout](https://tenstorrent.github.io/ttnn/latest/ttnn/tensor.html#layout).\n",
    "The current tile size is hard-coded to [32, 32]. It was determined to be the optimal size for a tile given the compute, memory and data transfer constraints.\n",
    "\n",
    "\n",
    "ttnn provides easy and intuitive way to convert from row-major layout to tile layout and back."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tensor in row-major layout:\n",
      "Shape Shape([3, 4]), padded shape: Shape([3, 4])\n",
      "Layout: Layout.ROW_MAJOR\n",
      "ttnn.Tensor([[ 0.04346,  0.51953,  0.17383,  0.08691],\n",
      "             [ 0.53906,  0.26953,  0.16602,  0.08301],\n",
      "             [ 0.97656,  0.48828,  0.91016,  0.95312]], shape=Shape([3, 4]), dtype=DataType::BFLOAT16, layout=Layout::ROW_MAJOR)\n",
      "Tensor in tile layout:\n",
      "Shape Shape([3, 4]), padded shape: Shape([32, 32])\n",
      "Layout: Layout.TILE\n",
      "ttnn.Tensor([[ 0.04346,  0.51953,  0.17383,  0.08691],\n",
      "             [ 0.53906,  0.26953,  0.16602,  0.08301],\n",
      "             [ 0.97656,  0.48828,  0.91016,  0.95312]], shape=Shape([3, 4]), dtype=DataType::BFLOAT16, layout=Layout::TILE)\n",
      "Tensor back in row-major layout:\n",
      "Shape Shape([3, 4]), padded shape: Shape([3, 4])\n",
      "Layout: Layout.ROW_MAJOR\n",
      "ttnn.Tensor([[ 0.04346,  0.51953,  0.17383,  0.08691],\n",
      "             [ 0.53906,  0.26953,  0.16602,  0.08301],\n",
      "             [ 0.97656,  0.48828,  0.91016,  0.95312]], shape=Shape([3, 4]), dtype=DataType::BFLOAT16, layout=Layout::ROW_MAJOR)\n"
     ]
    }
   ],
   "source": [
    "ttnn_tensor = ttnn.rand((3, 4), layout=ttnn.ROW_MAJOR_LAYOUT, device=device)\n",
    "print(f\"Tensor in row-major layout:\\nShape {ttnn_tensor.shape}, padded shape: {ttnn_tensor.padded_shape}\\nLayout: {ttnn_tensor.layout}\\n{ttnn_tensor}\")\n",
    "ttnn_tensor = ttnn.to_layout(ttnn_tensor, ttnn.TILE_LAYOUT)\n",
    "print(f\"Tensor in tile layout:\\nShape {ttnn_tensor.shape}, padded shape: {ttnn_tensor.padded_shape}\\nLayout: {ttnn_tensor.layout}\\n{ttnn_tensor}\")\n",
    "ttnn_tensor = ttnn.to_layout(ttnn_tensor, ttnn.ROW_MAJOR_LAYOUT)\n",
    "print(f\"Tensor back in row-major layout:\\nShape {ttnn_tensor.shape}, padded shape: {ttnn_tensor.padded_shape}\\nLayout: {ttnn_tensor.layout}\\n{ttnn_tensor}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that padding is automatically inserted to put the tensor into tile layout and it automatically removed after the tensor is converted back to row-major layout"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The conversion to tile layout can be done when caling `ttnn.from_torch`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tensor in row-major layout:\n",
      "Shape torch.Size([3, 4]); Layout: torch.strided\n",
      "Tensor in tile layout:\n",
      "Shape Shape([3, 4]); Layout: Layout.TILE\n"
     ]
    }
   ],
   "source": [
    "torch_tensor = ttnn.to_torch(ttnn.rand((3, 4), device=device, layout=ttnn.ROW_MAJOR_LAYOUT))\n",
    "ttnn_tensor = ttnn.from_torch(torch_tensor, layout=ttnn.TILE_LAYOUT)\n",
    "print(f\"Tensor in row-major layout:\\nShape {torch_tensor.shape}; Layout: {torch_tensor.layout}\")\n",
    "print(f\"Tensor in tile layout:\\nShape {ttnn_tensor.shape}; Layout: {ttnn_tensor.layout}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that `ttnn.to_torch` will always convert to row-major layout"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initialize tensors a and b with random values\n",
    "\n",
    "Create a tensor that can be used by a `ttnn` operation with the `ttnn.TILE_LAYOUT` and put the tensor on the `device`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input_tensor_a:\n",
      "ttnn.Tensor([[ 0.62109,  0.31055,  ...,  0.98828,  0.99219],\n",
      "             [ 0.69531,  0.34766,  ...,  0.72266,  0.36133],\n",
      "             ...,\n",
      "             [ 0.85156,  0.92578,  ...,  0.69531,  0.34766],\n",
      "             [ 0.58203,  0.78906,  ...,  0.29688,  0.64844]], shape=Shape([32, 32]), dtype=DataType::BFLOAT16, layout=Layout::TILE)\n",
      "input_tensor_b:\n",
      "ttnn.Tensor([[ 0.10596,  0.05298,  ...,  0.80078,  0.40039],\n",
      "             [ 0.69922,  0.84766,  ...,  0.94141,  0.47070],\n",
      "             ...,\n",
      "             [ 0.86719,  0.43359,  ...,  0.69922,  0.84766],\n",
      "             [ 0.06689,  0.53125,  ...,  0.28906,  0.14453]], shape=Shape([32, 32]), dtype=DataType::BFLOAT16, layout=Layout::TILE)\n"
     ]
    }
   ],
   "source": [
    "input_tensor_a = ttnn.rand((32, 32), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=device)\n",
    "input_tensor_b = ttnn.rand((32, 32), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=device)\n",
    "\n",
    "print(f\"input_tensor_a:\\n{input_tensor_a}\")\n",
    "print(f\"input_tensor_b:\\n{input_tensor_b}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Add tensor a and b\n",
    "\n",
    "`ttnn` supports operator overloading, therefore operator `+` can be used instead of `ttnn.add`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "output_tensor:\n",
      "ttnn.Tensor([[ 0.72656,  0.36328,  ...,  1.78906,  1.39062],\n",
      "             [ 1.39844,  1.19531,  ...,  1.66406,  0.83203],\n",
      "             ...,\n",
      "             [ 1.71875,  1.35938,  ...,  1.39844,  1.19531],\n",
      "             [ 0.64844,  1.32031,  ...,  0.58594,  0.79297]], shape=Shape([32, 32]), dtype=DataType::BFLOAT16, layout=Layout::TILE)\n"
     ]
    }
   ],
   "source": [
    "output_tensor = input_tensor_a + input_tensor_b\n",
    "print(f\"output_tensor:\\n{output_tensor}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inspect the output tensor of the add in ttnn\n",
    "\n",
    "As can be seen the tensor of the same shape, layout and dtype is produced"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "shape: Shape([32, 32])\n",
      "dtype: DataType.BFLOAT16\n",
      "layout: Layout.TILE\n"
     ]
    }
   ],
   "source": [
    "print(f\"shape: {output_tensor.shape}\")\n",
    "print(f\"dtype: {output_tensor.dtype}\")\n",
    "print(f\"layout: {output_tensor.layout}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In general we expect layout and dtype to stay the same when running most operations unless explicit arguments to modify them are passed in. However, there are obvious exceptions like an embedding operation that takes in `ttnn.uint32` and produces `ttnn.bfloat16`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Convert to torch and inspect the attributes of the torch tensor\n",
    "\n",
    "When converting the tensor to torch, `ttnn.to_torch` will move the tensor from the device, convert to tile layout and figure out the best data type to use on the torch side"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "shape: torch.Size([32, 32])\n",
      "dtype: torch.bfloat16\n"
     ]
    }
   ],
   "source": [
    "output_tensor = ttnn.to_torch(output_tensor)\n",
    "print(f\"shape: {output_tensor.shape}\")\n",
    "print(f\"dtype: {output_tensor.dtype}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Close the device\n",
    "\n",
    "Close the handle the device. This is a very important step as the device can hang currently if not closed properly"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2025-07-02 22:33:01.541 | info     |           Metal | Closing mesh device 1 (mesh_device.cpp:488)\n",
      "2025-07-02 22:33:01.543 | info     |           Metal | Closing mesh device 0 (mesh_device.cpp:488)\n",
      "2025-07-02 22:33:01.543 | info     |           Metal | Closing device 0 (device.cpp:469)\n",
      "2025-07-02 22:33:01.543 | info     |           Metal | Disabling and clearing program cache on device 0 (device.cpp:781)\n"
     ]
    }
   ],
   "source": [
    "ttnn.close_device(device)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
